{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58d596a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#planner\n",
    "from pydantic import BaseModel, Field\n",
    "from agents import Agent\n",
    "from datetime import datetime\n",
    "\n",
    "HOW_MANY_SEARCHES = 5\n",
    "\n",
    "INSTRUCTIONS = f\"You are a research assistant for a company. \\\n",
    "Given the company name, industry and a research query, come up with a set of web searches to perform to best answer the query considering areas such as \\\n",
    "user or customer feedback, competitor analytics, market trends, etc. The current date is {datetime.now().strftime('%Y-%m-%d')}. \\\n",
    "Output {HOW_MANY_SEARCHES} terms to query for.  Ensure the search queries are specific to the company.\"\n",
    "\n",
    "class WebSearchItem(BaseModel):\n",
    "    reason: str = Field(description=\"Your reasoning for why this search is important to the query.\")\n",
    "    query: str = Field(description=\"The search term to use for the web search.\")\n",
    "\n",
    "class WebSearchPlan(BaseModel):\n",
    "    searches: list[WebSearchItem] = Field(description=\"A list of web searches to perform to best answer the query.\")\n",
    "\n",
    "planner_agent = Agent(\n",
    "    name=\"PlannerAgent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=WebSearchPlan\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b900387e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#searcher\n",
    "from agents import Agent, WebSearchTool, ModelSettings\n",
    "\n",
    "INSTRUCTIONS = (\n",
    "    f\"You are a research assistant for a company. \\\n",
    "    Given a search term, you search the web for that term and produce a concise summary of the results. The summary must 2-3 paragraphs and less than 300 \"\n",
    "    \"words. Capture the main points. Write succintly, no need to have complete sentences or good \"\n",
    "    \"grammar. This will be consumed by someone synthesizing a report, so its vital you capture the \"\n",
    "    \"essence and ignore any fluff. Do not include any additional commentary other than the summary itself.\"\n",
    ")\n",
    "\n",
    "search_agent = Agent(\n",
    "    name=\"Search agent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    tools=[WebSearchTool(search_context_size=\"low\")],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    model_settings=ModelSettings(tool_choice=\"required\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1588c31f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#writer\n",
    "from pydantic import BaseModel, Field\n",
    "from agents import Agent\n",
    "\n",
    "INSTRUCTIONS = (\n",
    "    f\"You are a senior researcher for a company.  You are tasked with writing a cohesive report for a research query. \\\n",
    "    You will be provided with the company name, industry, original query, and some initial research done by a research assistant.\\n\"\n",
    "    \"You should first come up with an outline for the report that describes the structure and \"\n",
    "    \"flow of the report. Then, generate the report and return that as your final output.\\n\"\n",
    "    \"The final output should be in markdown format, and it should be lengthy and detailed. Aim \"\n",
    "    \"for 5-10 pages of content, at least 1000 words.\"\n",
    ")\n",
    "\n",
    "class ReportData(BaseModel):\n",
    "    short_summary: str = Field(description=\"A short 2-3 sentence summary of the findings.\")\n",
    "\n",
    "    markdown_report: str = Field(description=\"The final report\")\n",
    "\n",
    "    follow_up_questions: list[str] = Field(description=\"Suggested topics to research further\")\n",
    "\n",
    "\n",
    "writer_agent = Agent(\n",
    "    name=\"WriterAgent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ReportData\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3b4e259",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Dict\n",
    "import smtplib\n",
    "from email.mime.text import MIMEText\n",
    "from agents import Agent, function_tool\n",
    "\n",
    "@function_tool\n",
    "def send_email(subject: str, html_body: str, email_address: str) -> Dict[str, str]:\n",
    "    \"\"\" Send out an email with the given subject and HTML body to the user \"\"\"\n",
    "\n",
    "    # Your Gmail credentials\n",
    "    gmail_user = os.getenv(\"FROM_EMAIL\") \n",
    "    gmail_app_password = os.getenv(\"GOOGLE_APP_PW\")\n",
    "    to = email_address #os.getenv(\"TO_EMAIL\") \n",
    "\n",
    "    if not all([gmail_user, gmail_app_password, to]):\n",
    "        raise ValueError(\"Missing one or more required environment variables: FROM_EMAIL, GOOGLE_APP_PW, TO_EMAIL\")\n",
    "\n",
    "    # Create the email\n",
    "    msg = MIMEText(html_body, \"html\")\n",
    "    msg['Subject'] = subject or \"\"\n",
    "    msg['From'] = gmail_user or \"\"\n",
    "    msg['To'] = to or \"\"\n",
    "\n",
    "    # Send the email\n",
    "    try:\n",
    "        with smtplib.SMTP('smtp.gmail.com', 587) as server:\n",
    "            server.starttls()  # Secure the connection\n",
    "            server.login(gmail_user, gmail_app_password)\n",
    "            server.send_message(msg)\n",
    "        print('Email sent successfully!')\n",
    "        return {\"status\": \"success\"}\n",
    "    except Exception as e:\n",
    "        print(f'Failed to send email: {e}')\n",
    "\n",
    "\n",
    "#handoff prompt recommended by openai\n",
    "RECOMMENDED_PROMPT_PREFIX = \"# System context\\nYou are part of a multi-agent system called the Agents SDK, designed to make agent coordination and execution easy. \\\n",
    "Agents uses two primary abstraction: **Agents** and **Handoffs**. An agent encompasses instructions and tools and can hand off a conversation to another agent when \\\n",
    "appropriate. Handoffs are achieved by calling a handoff function, generally named `transfer_to_<agent_name>`. Transfers between agents are handled seamlessly in the background;\\\n",
    "do not mention or draw attention to these transfers in your conversation with the user.\\n\"\n",
    "\n",
    "INSTRUCTIONS = f\"\"\"{RECOMMENDED_PROMPT_PREFIX} You are able to send a nicely formatted HTML email based on a detailed report.\n",
    "You will be provided with a detailed report and the email address to which it should be sent. You should use your tool to send one email, providing the \n",
    "report converted into clean, well presented HTML with an appropriate subject line.\"\"\"\n",
    "\n",
    "email_agent = Agent(\n",
    "    name=\"Email agent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    tools=[send_email],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    handoff_description=\"Convert a report to HTML and send it\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c06c9f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#research manager\n",
    "from agents import Agent, Runner, trace, gen_trace_id, SQLiteSession, handoff, trace\n",
    "from agents.extensions import handoff_filters\n",
    "\n",
    "\n",
    "# Convert agents to tools\n",
    "planner_tool = planner_agent.as_tool(tool_name=\"planner_agent\", tool_description=\"Create search strategy\")\n",
    "search_tool = search_agent.as_tool(tool_name=\"search_agent\", tool_description=\"Execute web searches and summarises results\")\n",
    "writer_tool = writer_agent.as_tool(tool_name=\"writer_agent\", tool_description=\"Generate research report\")\n",
    "\n",
    "# Research Manager Agent\n",
    "\n",
    "#handoff prompt recommended by openai\n",
    "RECOMMENDED_PROMPT_PREFIX = \"# System context\\nYou are part of a multi-agent system called the Agents SDK, designed to make agent coordination and execution easy. \\\n",
    "Agents uses two primary abstraction: **Agents** and **Handoffs**. An agent encompasses instructions and tools and can hand off a conversation to another agent when \\\n",
    "appropriate. Handoffs are achieved by calling a handoff function, generally named `transfer_to_<agent_name>`. Transfers between agents are handled seamlessly in the background;\\\n",
    "do not mention or draw attention to these transfers in your conversation with the user.\\n\"\n",
    "\n",
    "INSTRUCTIONS = (\n",
    "    f\"\"\"{RECOMMENDED_PROMPT_PREFIX} You research companies and create reports. You have these tools:\n",
    "- planner_tool: Create search strategies\n",
    "- search_tool: Find information  \n",
    "- writer_tool: Create reports\n",
    "\n",
    "Look at the conversation history to understand what's needed:\n",
    "- New research request? Plan, search, write.\n",
    "- Feedback on existing research? Update appropriately.\n",
    "- Email request? Transfer to email agent.\n",
    "\n",
    "Use your intelligence to decide which tools to use and when.\n",
    "\n",
    "Always output only the final markdown report - no commentary or explanations.\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "#simplify handoff\n",
    "email_handoff = handoff(\n",
    "    agent=email_agent,\n",
    "    input_filter=handoff_filters.remove_all_tools,  # This removes all tool calls from history\n",
    "    tool_description_override=\"Send the research report via email to the specified email address\"\n",
    ")\n",
    "\n",
    "research_manager_agent = Agent(\n",
    "    name=\"Research Manager\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    model=\"gpt-4o\",\n",
    "    handoffs=[email_handoff],\n",
    "    tools=[planner_tool, search_tool, writer_tool]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5ae4c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "#research function\n",
    "\n",
    "session_store = {}\n",
    "\n",
    "async def run_research(company: str, industry: str, query: str, feedback: str, email_trigger: str):\n",
    "    \n",
    "    #session handling - to ensure session persists    \n",
    "    session_id = f\"research_{company}_{query}\"\n",
    "    if session_id not in session_store:\n",
    "        session_store[session_id] = SQLiteSession(session_id)\n",
    "    session = session_store[session_id] \n",
    "  \n",
    "    if hasattr(session, 'trace_id') and session.trace_id:\n",
    "        trace_id = session.trace_id\n",
    "    else:\n",
    "        trace_id = gen_trace_id()\n",
    "        session.trace_id = trace_id\n",
    "    \n",
    "    with trace(\"Research trace\", trace_id=trace_id):\n",
    "        print(f\"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}\")\n",
    "\n",
    "        if not feedback and not email_trigger:\n",
    "            # Initial research\n",
    "            yield \"Researching...\"\n",
    "            \n",
    "            result = await Runner.run(\n",
    "                research_manager_agent, \n",
    "                f\"Research: Company: {company} | Industry: {industry} | Query: {query}\",\n",
    "                session=session\n",
    "            )\n",
    "            yield result.final_output\n",
    "        \n",
    "        elif feedback and not email_trigger:\n",
    "            # Feedback processing\n",
    "            yield \"Processing feedback...\"\n",
    "            \n",
    "            result = await Runner.run(\n",
    "                research_manager_agent, \n",
    "                f\"Based on your previous research, here is user feedback: {feedback}\\n\\n\\\n",
    "                Please update and improve the existing research report based on this feedback. Do not start over - build upon what you already provided.\",\n",
    "                session=session\n",
    "            )\n",
    "            yield result.final_output\n",
    "\n",
    "        elif email_trigger.startswith(\"EMAIL_REPORT\"):\n",
    "            \n",
    "            yield \"Preparing to email report...\"\n",
    "            \n",
    "            email_address = email_trigger.split(\"Email address: \")[1]\n",
    "            result = await Runner.run(\n",
    "                research_manager_agent,\n",
    "                f\"The user wants to email the report. Please hand off to the emailer agent to send the final research report and share the email: {email_address}\",\n",
    "                session=session\n",
    "            )\n",
    "            yield result.final_output\n",
    "\n",
    "        else:\n",
    "            yield \"The research process is complete.  Please clear and start again.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "996b128f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# functions for gradio\n",
    "import gradio as gr\n",
    "import re\n",
    "\n",
    "#email functions\n",
    "def trigger_email(email):\n",
    "    return f\"EMAIL_REPORT, Email address: {email}\"\n",
    "\n",
    "def show_email_fields():\n",
    "    return gr.update(visible=True,interactive=True), gr.update(visible=True), gr.update(visible=True,interactive=True)\n",
    "\n",
    "def validate_email(email):\n",
    "    \"Validates an email address using a regular expression.\"\n",
    "    \n",
    "    if not email:\n",
    "        return \"Email cannot be empty.\"\n",
    "    pattern = r\"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$\"\n",
    "    if re.match(pattern, email):\n",
    "        return \"Valid email address.\"\n",
    "    else:\n",
    "        return \"Invalid email address format.\"\n",
    "\n",
    "#main function\n",
    "#first run\n",
    "async def run_initial(company: str, industry: str, query: str):\n",
    "\n",
    "    # Validation\n",
    "    if not company or not company.strip():\n",
    "        raise gr.Error(\"Please enter a company name\")\n",
    "    if not industry or not industry.strip():\n",
    "        raise gr.Error(\"Please enter an industry\")\n",
    "    if not query or not query.strip():\n",
    "        raise gr.Error(\"Please enter a research query\")\n",
    "\n",
    "    async for chunk in run_research(company, industry, query, None, None):\n",
    "        yield chunk\n",
    "\n",
    "#feedback run\n",
    "async def run_feedback(company: str, industry: str, query: str, feedback: str):\n",
    "    # Validation\n",
    "    if not feedback or not feedback.strip():\n",
    "        raise gr.Error(\"Please enter your feedback\")\n",
    "\n",
    "    async for chunk in run_research(company, industry, query, feedback, None):\n",
    "        yield chunk\n",
    "\n",
    "#email handoff run\n",
    "async def run_email(company: str, industry: str, query: str, feedback: str, email_trigger: str):\n",
    "    async for chunk in run_research(company, industry, query, feedback, email_trigger):\n",
    "        yield chunk\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46cbe803",
   "metadata": {},
   "outputs": [],
   "source": [
    "# final run and gradio\n",
    "import gradio as gr\n",
    "\n",
    "with gr.Blocks(theme=gr.themes.Ocean()) as ui:\n",
    "    gr.Markdown(\"Deep Research for Companies\")\n",
    "    gr.Markdown(\"Generate comprehensive company research reports with AI-powered analysis\")\n",
    "    \n",
    "    with gr.Row():\n",
    "        with gr.Column(scale=1):\n",
    "            org_textbox = gr.Textbox(label=\"Organisation Name\",placeholder=\"e.g. ComplAI\",info=\"The company you want to research\")\n",
    "            industry_textbox = gr.Textbox(label=\"Industry\", placeholder=\"e.g. Finance, Healthcare, SaaS\",info=\"Primary industry or sector\")\n",
    "            query_textbox = gr.Textbox(label=\"Research Topic\",placeholder=\"e.g. Analyze competitor pricing strategies\",lines=3,info=\"What specific aspect would you like to research?\")\n",
    "            \n",
    "            with gr.Row():\n",
    "                run_button = gr.Button(\"Start Research\", variant=\"primary\", size=\"lg\")\n",
    "\n",
    "        with gr.Column(scale=1):\n",
    "            feedback_textbox = gr.Textbox(label=\"Feedback\", placeholder=\"e.g. Can you provide more information on...?\",lines=3,info=\"Provide feedback to improve the research\")\n",
    "            email_trigger_textbox = gr.Textbox(label=\"Email trigger\", visible=False)\n",
    "            feedback_button = gr.Button(\"Provide Feedback\", variant=\"primary\",scale=1)\n",
    "            send_button = gr.Button(\"Email me the report\", variant =\"secondary\",scale=1)\n",
    "            with gr.Column(scale=1):\n",
    "                email_address = gr.Textbox(label=\"Email address:\",type=\"email\",visible=False,interactive=True)\n",
    "                email_validation = gr.Markdown(visible=False)\n",
    "                confirm_send_button = gr.Button(\"Go, please send!\", variant =\"primary\",visible=False)\n",
    "        \n",
    "        with gr.Column(scale=2):\n",
    "            report = gr.Markdown(label=\"Research Report\",value=\"Your research report will appear here...\",height=600)\n",
    "    \n",
    "    # Event handlers\n",
    "    run_button.click(fn=run_initial, inputs=[org_textbox, industry_textbox, query_textbox], outputs=report, show_progress=True)\n",
    "    feedback_button.click(fn=run_feedback, inputs=[org_textbox, industry_textbox, query_textbox, feedback_textbox], outputs=report, show_progress=True)\n",
    "    send_button.click(fn=show_email_fields,inputs=[],outputs=[email_address, email_validation, confirm_send_button])\n",
    "    confirm_send_button.click(fn=validate_email, inputs=email_address, outputs=email_validation).then(fn=trigger_email, inputs=email_address, outputs=email_trigger_textbox).then(\n",
    "    fn=run_email, inputs=[org_textbox, industry_textbox, query_textbox, feedback_textbox, email_trigger_textbox], outputs=report)\n",
    "        \n",
    "    gr.ClearButton(components=[org_textbox,industry_textbox,query_textbox,feedback_textbox,report,email_address],value=\"Clear all fields\")\n",
    "\n",
    "ui.launch(inbrowser=True, share=False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edd5f581",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67869a9a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86540e00",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
